# 使用 ARG 指令接受 Pytorch 镜像版本
ARG IMAGE_VERSION
ARG ROOT_PASSWORD

# 使用官方 pytorch 基础镜像
FROM pytorch/pytorch:${IMAGE_VERSION}

# 避免在安装过程中由于缺少交互导致的挂起
ENV DEBIAN_FRONTEND=noninteractive

# 安装 SSH Server/nano/net-tools
RUN sed -i 's/http:\/\/archive.ubuntu.com/https:\/\/mirror.sjtu.edu.cn/g' /etc/apt/sources.list
    
RUN apt-get update && \
    apt-get install -y openssh-server nano net-tools inetutils-ping && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    mkdir /var/run/sshd && \
    mkdir -p /root/.ssh && \
    touch /root/.ssh/authorized_keys

# 使用环境变量设置 root 密码
ARG ROOT_PASSWORD
RUN echo "root:${ROOT_PASSWORD}" | chpasswd

# 配置 SSH 允许 root 登录并使用公钥认证
RUN sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config \
    && sed -i 's/#PubkeyAuthentication yes/PubkeyAuthentication yes/' /etc/ssh/sshd_config \
    && sed -i 's@#AuthorizedKeysFile.*@AuthorizedKeysFile .ssh/authorized_keys@' /etc/ssh/sshd_config

# 复制公钥文件到容器中
COPY keys/ /root/.ssh/keys/

# 将所有公钥文件写入到 authorized_keys 中
RUN cat /root/.ssh/keys/* >> /root/.ssh/authorized_keys \
    && chmod 600 /root/.ssh/authorized_keys \
    && chmod 700 ~/.ssh

RUN echo 'export PATH=/opt/conda/bin:$PATH' >> ~/.bashrc \
    && . ~/.bashrc \
    && conda init


# 放开端口 22
EXPOSE 22

# 容器启动时运行 SSH 服务
CMD ["/usr/sbin/sshd", "-D"]
